from trace_visualizer_py.common import common_data

from PyQt5 import QtWidgets
from PyQt5.QtCore import QTimer

import numpy as np
# from PyQt5.QtOpenGL import QGLWidget
from OpenGL.GL import *
from OpenGL.GLU import *

class OpenGLWidget(QtWidgets.QOpenGLWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setMinimumSize(200, 200)
        # 物体姿态 (位置和四元数旋转)
        self.position = np.array([0.0, 0.0, 0.0])  # x, y, z
        self.orientation = np.array([1.0, 0.0, 0.0, 0.0])  # 四元数 (w, x, y, z)
        
        # 坐标系大小
        self.axis_length = 1.0
        
        # 设置定时器用于动画演示
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_pose_demo)
        self.timer.start(50)  # 20 fps
        
    def initializeGL(self):
        glEnable(GL_DEPTH_TEST)
        glClearColor(0.0, 0.0, 0.0, 1.0)
        
    def resizeGL(self, w, h):
        glViewport(0, 0, w, h)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(45, w/h, 0.1, 100.0)
        glMatrixMode(GL_MODELVIEW)
        
    def paintGL(self):
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glLoadIdentity()
        
        # # 设置视角
        # gluLookAt(0, 0, 5, 0, 0, 0, 0, 1, 0)
        
        # # 绘制一个彩色立方体
        # glBegin(GL_QUADS)
        
        # # 前面
        # glColor3f(1.0, 0.0, 0.0)
        # glVertex3f(-1.0, -1.0, 1.0)
        # glVertex3f(1.0, -1.0, 1.0)
        # glVertex3f(1.0, 1.0, 1.0)
        # glVertex3f(-1.0, 1.0, 1.0)
        
        # # 后面
        # glColor3f(0.0, 1.0, 0.0)
        # glVertex3f(-1.0, -1.0, -1.0)
        # glVertex3f(-1.0, 1.0, -1.0)
        # glVertex3f(1.0, 1.0, -1.0)
        # glVertex3f(1.0, -1.0, -1.0)
        
        # # 省略其他面的绘制...
        
        # glEnd()
        
        # 设置视角
        gluLookAt(3, 3, 3, 0, 0, 0, 0, 1, 0)
        
        # 绘制全局坐标系
        self.draw_global_axes()
        
        # 应用物体姿态变换
        glPushMatrix()
        self.apply_pose_transform()
        
        # 绘制物体局部坐标系
        self.draw_local_axes()
        
        # 绘制物体 (这里用一个立方体表示)
        self.draw_cube()
        
        glPopMatrix()

    def apply_pose_transform(self):
        """应用位置和旋转变换"""
        # 应用位置
        glTranslatef(self.position[0], self.position[1], self.position[2])
        
        # 将四元数转换为旋转矩阵并应用
        w, x, y, z = self.orientation
        angle = 2 * np.arccos(w) * 180 / np.pi
        norm = np.sqrt(x*x + y*y + z*z)
        if norm > 0:
            x /= norm
            y /= norm
            z /= norm
            glRotatef(angle, x, y, z)
    
    def draw_global_axes(self):
        """绘制全局坐标系"""
        glBegin(GL_LINES)
        # X轴 (红色)
        glColor3f(1.0, 0.0, 0.0)
        glVertex3f(0.0, 0.0, 0.0)
        glVertex3f(self.axis_length, 0.0, 0.0)
        # Y轴 (绿色)
        glColor3f(0.0, 1.0, 0.0)
        glVertex3f(0.0, 0.0, 0.0)
        glVertex3f(0.0, self.axis_length, 0.0)
        # Z轴 (蓝色)
        glColor3f(0.0, 0.0, 1.0)
        glVertex3f(0.0, 0.0, 0.0)
        glVertex3f(0.0, 0.0, self.axis_length)
        glEnd()
    
    def draw_local_axes(self):
        """绘制物体局部坐标系"""
        axis_length = self.axis_length * 0.5  # 局部坐标系比全局的小一些
        
        glBegin(GL_LINES)
        # X轴 (浅红色)
        glColor3f(1.0, 0.5, 0.5)
        glVertex3f(0.0, 0.0, 0.0)
        glVertex3f(axis_length, 0.0, 0.0)
        # Y轴 (浅绿色)
        glColor3f(0.5, 1.0, 0.5)
        glVertex3f(0.0, 0.0, 0.0)
        glVertex3f(0.0, axis_length, 0.0)
        # Z轴 (浅蓝色)
        glColor3f(0.5, 0.5, 1.0)
        glVertex3f(0.0, 0.0, 0.0)
        glVertex3f(0.0, 0.0, axis_length)
        glEnd()
    
    def draw_cube(self):
        """绘制一个简单的立方体表示物体"""
        size = 0.3
        glBegin(GL_QUADS)
        
        # 前面 (红色)
        glColor3f(1.0, 0.0, 0.0)
        glVertex3f(-size, -size, size)
        glVertex3f(size, -size, size)
        glVertex3f(size, size, size)
        glVertex3f(-size, size, size)
        
        # 后面 (绿色)
        glColor3f(0.0, 1.0, 0.0)
        glVertex3f(-size, -size, -size)
        glVertex3f(-size, size, -size)
        glVertex3f(size, size, -size)
        glVertex3f(size, -size, -size)
        
        # 其他面省略...
        
        glEnd()

    def update_pose(self, position, orientation):
        """从外部更新姿态数据"""
        self.position = position  # [x, y, z]
        self.orientation = orientation  # [w, x, y, z] 四元数
        self.update()  # 触发重绘

    def update_pose_demo(self):
        """更新姿态用于演示"""
        # 这里只是一个演示，实际应用中应该从你的数据源获取姿态
        import math
        import time
        
        t = time.time()
        
        # 更新位置 (圆周运动)
        # radius = 1.0
        # self.position[0] = radius * math.sin(t)
        # self.position[1] = 0.5 * math.cos(t)
        # self.position[2] = radius * math.cos(t)
        # self.position = 0
        
        # 更新旋转 (随时间旋转)
        angle = t * 45  # 45度/秒
        half_angle = math.radians(angle) / 2
        self.orientation[0] = math.cos(half_angle)  # w
        # self.orientation[1] = math.sin(half_angle)  # x
        self.orientation[2] = 1  # y
        self.orientation[3] = 1 # math.sin(half_angle)  # z
        
        self.update()  # 触发重绘

common_data['gl_widget'] = OpenGLWidget()

# class MainWindow(QtWidgets.QMainWindow):
#     def __init__(self):
#         super().__init__()
        
#         # 创建主窗口的其他部件...
        
#         # 添加3D显示部件
#         self.threeDWidget = OpenGLWidget()  # 或 OpenGLWidget() 或 VTKWidget()
        
#         # 将3D部件添加到布局中
#         mainLayout = QtWidgets.QHBoxLayout()
#         # mainLayout.addWidget(existingWidget)  # 你已有的部件
#         mainLayout.addWidget(self.threeDWidget)
        
#         centralWidget = QtWidgets.QWidget()
#         centralWidget.setLayout(mainLayout)
#         self.setCentralWidget(centralWidget)